#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Agent 扩展模块 - 演示如何扩展 Agent 功能
包括数据库支持、知识库集成等
"""

import os
import json
import sqlite3
from typing import Optional, Dict, Any, List
from datetime import datetime

# LangChain imports
try:
    from langchain.agents import Tool
    from langchain.schema import Document
except ImportError:
    print("❌ 请先安装 LangChain: pip install langchain")


class DatabaseTool:
    """数据库工具 - 支持简单的 SQLite 操作"""
    
    def __init__(self, db_path: str = "agent_data.db"):
        self.db_path = db_path
        self.init_database()
    
    def init_database(self):
        """初始化数据库"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            # 创建用户记录表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS user_records (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    user_id TEXT,
                    content TEXT,
                    category TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # 创建对话历史表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS chat_history (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    session_id TEXT,
                    user_message TEXT,
                    agent_response TEXT,
                    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            conn.commit()
            conn.close()
            print("✅ 数据库初始化完成")
            
        except Exception as e:
            print(f"❌ 数据库初始化失败: {e}")
    
    def save_record(self, query: str) -> str:
        """保存用户记录"""
        try:
            # 解析查询 - 简单格式: "保存记录: 用户ID|内容|类别"
            if not query.startswith("保存记录:"):
                return "❌ 格式错误。正确格式: 保存记录: 用户ID|内容|类别"
            
            data = query.replace("保存记录:", "").strip()
            parts = data.split("|")
            
            if len(parts) != 3:
                return "❌ 参数错误。需要: 用户ID|内容|类别"
            
            user_id, content, category = [p.strip() for p in parts]
            
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('''
                INSERT INTO user_records (user_id, content, category)
                VALUES (?, ?, ?)
            ''', (user_id, content, category))
            
            conn.commit()
            record_id = cursor.lastrowid
            conn.close()
            
            return f"✅ 记录已保存，ID: {record_id}"
            
        except Exception as e:
            return f"❌ 保存失败: {e}"
    
    def query_records(self, query: str) -> str:
        """查询用户记录"""
        try:
            # 解析查询 - 简单格式: "查询记录: 用户ID" 或 "查询记录: 用户ID|类别"
            if not query.startswith("查询记录:"):
                return "❌ 格式错误。正确格式: 查询记录: 用户ID 或 查询记录: 用户ID|类别"
            
            data = query.replace("查询记录:", "").strip()
            parts = data.split("|")
            
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            if len(parts) == 1:
                # 查询用户的所有记录
                user_id = parts[0].strip()
                cursor.execute('''
                    SELECT id, content, category, created_at 
                    FROM user_records 
                    WHERE user_id = ?
                    ORDER BY created_at DESC
                ''', (user_id,))
            else:
                # 查询特定类别的记录
                user_id, category = [p.strip() for p in parts]
                cursor.execute('''
                    SELECT id, content, category, created_at 
                    FROM user_records 
                    WHERE user_id = ? AND category = ?
                    ORDER BY created_at DESC
                ''', (user_id, category))
            
            records = cursor.fetchall()
            conn.close()
            
            if not records:
                return "📝 未找到相关记录"
            
            result = f"📋 找到 {len(records)} 条记录:\n"
            for record in records[:10]:  # 最多显示10条
                record_id, content, category, created_at = record
                result += f"   [{record_id}] {category}: {content} ({created_at})\n"
            
            if len(records) > 10:
                result += f"   ... 还有 {len(records) - 10} 条记录"
            
            return result.strip()
            
        except Exception as e:
            return f"❌ 查询失败: {e}"


class KnowledgeBase:
    """简单的知识库实现"""
    
    def __init__(self, kb_path: str = "knowledge_base.json"):
        self.kb_path = kb_path
        self.knowledge = self.load_knowledge()
    
    def load_knowledge(self) -> Dict[str, Any]:
        """加载知识库"""
        if os.path.exists(self.kb_path):
            try:
                with open(self.kb_path, 'r', encoding='utf-8') as f:
                    return json.load(f)
            except Exception as e:
                print(f"❌ 加载知识库失败: {e}")
                return {}
        return {}
    
    def save_knowledge(self):
        """保存知识库"""
        try:
            with open(self.kb_path, 'w', encoding='utf-8') as f:
                json.dump(self.knowledge, f, ensure_ascii=False, indent=2)
        except Exception as e:
            print(f"❌ 保存知识库失败: {e}")
    
    def add_knowledge(self, query: str) -> str:
        """添加知识"""
        try:
            # 格式: "添加知识: 关键词|内容"
            if not query.startswith("添加知识:"):
                return "❌ 格式错误。正确格式: 添加知识: 关键词|内容"
            
            data = query.replace("添加知识:", "").strip()
            parts = data.split("|", 1)  # 只分割一次
            
            if len(parts) != 2:
                return "❌ 参数错误。需要: 关键词|内容"
            
            keyword, content = [p.strip() for p in parts]
            
            if keyword not in self.knowledge:
                self.knowledge[keyword] = []
            
            self.knowledge[keyword].append({
                "content": content,
                "timestamp": datetime.now().isoformat()
            })
            
            self.save_knowledge()
            return f"✅ 知识已添加到关键词 '{keyword}'"
            
        except Exception as e:
            return f"❌ 添加知识失败: {e}"
    
    def search_knowledge(self, query: str) -> str:
        """搜索知识"""
        try:
            # 格式: "搜索知识: 关键词"
            if not query.startswith("搜索知识:"):
                return "❌ 格式错误。正确格式: 搜索知识: 关键词"
            
            keyword = query.replace("搜索知识:", "").strip()
            
            # 精确匹配
            if keyword in self.knowledge:
                items = self.knowledge[keyword]
                result = f"🔍 找到关键词 '{keyword}' 的 {len(items)} 条知识:\n"
                for i, item in enumerate(items[-5:], 1):  # 最多显示最新5条
                    result += f"   {i}. {item['content']}\n"
                return result.strip()
            
            # 模糊搜索
            matches = []
            for key, items in self.knowledge.items():
                if keyword.lower() in key.lower():
                    matches.extend([(key, item) for item in items])
            
            if matches:
                result = f"🔍 模糊搜索到 {len(matches)} 条相关知识:\n"
                for i, (key, item) in enumerate(matches[-5:], 1):
                    result += f"   {i}. [{key}] {item['content']}\n"
                return result.strip()
            
            return f"❌ 未找到关键词 '{keyword}' 相关的知识"
            
        except Exception as e:
            return f"❌ 搜索知识失败: {e}"


def create_extension_tools(db_path: str = "agent_data.db", kb_path: str = "knowledge_base.json") -> List[Tool]:
    """创建扩展工具列表"""
    
    # 数据库工具
    db_tool = DatabaseTool(db_path)
    
    # 知识库工具
    kb_tool = KnowledgeBase(kb_path)
    
    return [
        Tool(
            name="保存记录",
            description="保存用户记录到数据库。格式: 保存记录: 用户ID|内容|类别",
            func=db_tool.save_record
        ),
        Tool(
            name="查询记录", 
            description="从数据库查询用户记录。格式: 查询记录: 用户ID 或 查询记录: 用户ID|类别",
            func=db_tool.query_records
        ),
        Tool(
            name="添加知识",
            description="添加知识到知识库。格式: 添加知识: 关键词|内容",
            func=kb_tool.add_knowledge
        ),
        Tool(
            name="搜索知识",
            description="从知识库搜索知识。格式: 搜索知识: 关键词",
            func=kb_tool.search_knowledge
        )
    ]


def extend_agent_example():
    """扩展 Agent 的示例"""
    print("🔧 Agent 扩展示例")
    print("=" * 50)
    
    # 创建扩展工具
    extension_tools = create_extension_tools()
    
    print("✅ 可用的扩展工具:")
    for tool in extension_tools:
        print(f"   📋 {tool.name}: {tool.description}")
    
    print("\n💡 使用示例:")
    print("   保存记录: user123|今天学习了Python|学习")
    print("   查询记录: user123")
    print("   添加知识: Python|Python是一种高级编程语言")
    print("   搜索知识: Python")


if __name__ == "__main__":
    extend_agent_example()
